import sys
import os
from json import dumps
import numpy as np
import pandas as pd

file_path = os.path.abspath(__file__)
dir_path = os.path.dirname(file_path)
os.chdir(dir_path)

# 加载微生物与疾病列表
dd = np.loadtxt('disease.txt', delimiter='\t', dtype=[('id', 'i4'), ('disease', 'S50')])
dm = np.loadtxt('microbe.txt', delimiter='\t', dtype=[('id', 'i4'), ('micro', 'S50')])
# 已经排名的疾病-微生物的关联度矩阵
SR = np.load('SortedResult.npy')
# 疾病-序号映射表
name_id = {v[1].decode():v[0]-1 for v in dd}

def pridect(disease: str):
    idx = name_id[disease]
    data = pd.DataFrame({'Rank':(range(1, 11)),
                         'microbe': (map(lambda x: dm[int(x[0])][1].decode(),
                                         SR[idx][:10]))})
    data.set_index('Rank', inplace=True)
    ans = [(0, disease)]
    ans += [(idx, m['microbe']) for idx, m in data.iterrows()]
    return dumps(ans)

if __name__ == '__main__':
    print(pridect(sys.argv[1]))